"""
用保存的 score_model ckpt 替换 AltDiffusion unet 中的 diffusion_pytorch_model.bin 文件
新创建 EulerDummyScheduler , 实现用于 reflow 模型的最简单 euler 采样方法
如此采样可使用和 AltDiffusionPipeline 兼容的方法

本 script 为产生支持指标计算的 sample 设计
需要提前准备好 captions ; 暂定按行存储在 .txt 文件中
支持的功能：
- 指定产生文件的根目录
- 数据产生格式支持3种指标的计算
- IS 和 FID ：根目录下 images 文件夹存放所有图片, index.json 存放 args 信息
- CLIPSim ：根目录下 images 存放图片 , captions.json 存放 index.json 存放 args 信息
- - images 下文件夹编号为 1,2,3... , 每个文件夹下固定 y 张图片 (clip候选图片) , 从中计算挑选 CLIPSim 最高的作为这一组的最终分数
- - captions.json 中为 {'1': caption1 , '2': caption2 , ...} 的格式

超参数：
- 随机种子 ; 
- 目标文件夹位置 ; 
- devices ; 
- batch size ; 
- score model ckpt 的路径 ; (在代码中进行文件移动和替换)
- 样本数量(按照文本数量计算) ; 
- 推理步数 ; 
- guidance scale ; 
- 指标种类 (IS , FID , CLIPSim) (指标种类决定数据保存的格式)
- - 如果是 CLIPSim 还需要保存 候选图片数量
"""
import sys
sys.path.append('.')
use_controlnet=False
if use_controlnet:
    sys.path.insert(0, "src/diffusers-controlnet") # 确保从 diffusers-controlnet 版本加载 model 和 pipeline

from pathlib import Path
import torch
from tqdm.auto import trange
from loguru import logger
import json
import multiprocessing as mp
import os
import time
from diffusers import UNet2DConditionModel
from argparse import Namespace, ArgumentParser, ArgumentTypeError


class HybridModel(torch.nn.Module):
    def __init__(self, zt_compressed_model, base_model, zt_compress_rate) -> None:
        super().__init__()
        self.zt_compressed_model = zt_compressed_model
        self.base_model = base_model

        self.zt_compress_rate = zt_compress_rate
        self.eps = 1e-3

    def forward(self, sample, timestep, **condition):
        if timestep < (999*(self.zt_compress_rate+self.eps)):
            timestep = timestep / self.zt_compress_rate
            model = self.zt_compressed_model
        else:
            model = self.base_model
        return model(sample, timestep=timestep, **condition)

    @property
    def device(self):
        return self.base_model.device
    
    def modules(self):
        return self.base_model.modules()
    
    @property
    def config(self):
        return self.base_model.config
    
    @property
    def in_channels(self):
        return self.base_model.in_channels
    
    @property
    def dtype(self):
        return self.base_model.dtype
    
    


from reflow.utils import set_seed, nothing, _PIPELINES, _SCHEDULERS, devide_to_groups
if use_controlnet:
    from diffusers import AltDiffusionControlNetPipeline, StableDiffusionControlNetPipeline
    _PIPELINES['stable_diffusion_controlnet']=StableDiffusionControlNetPipeline
    _PIPELINES['alt_diffusion_controlnet']=AltDiffusionControlNetPipeline

def produce(caps, idx_se, device, args, save_dir, image_dir: Path, ):
    set_seed(args.seed)
    logger.add(str(save_dir / 'run.log'))
    
    pipeline_cls = _PIPELINES[args.pipeline]
    scheduler_cls = _SCHEDULERS[args.scheduler]
    dtype=torch.float16
    pipe = pipeline_cls.from_pretrained(
        args.pipeline_ckpt,
        torch_dtype=dtype,
        safety_checker=None,
        requires_safety_checker=False,
    )

    pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config)
    if args.scheduler == 'euler_dummy':
        pipe.scheduler.overide_timesteps = args.overide_timesteps
        
    # 替换 ckpt
    if not nothing(args.ckpt_path):
        pipe.unet.load_state_dict(torch.load(args.ckpt_path, map_location='cpu'), strict=True)
        logger.info(
            f'use ckpt <<{args.ckpt_path}>>')
    else:
        logger.warning(f'not specify pretrained score model path')

    if args.load_lora:
        pipe.unet.load_attn_procs(args.lora_weights_path)
        
    if args.load_controlnet and use_controlnet:
        pipe.controlnet.load_state_dict(torch.load(args.controlnet_path, map_location='cpu'), strict=True)
        
    if args.load_zt_compressed_model:
        assert not nothing(args.ckpt_path), "need base model!"
        zt_compressed_model = UNet2DConditionModel.from_pretrained(
            args.pipeline_ckpt,
            subfolder='unet',
        )
        zt_compressed_model.load_state_dict(torch.load(args.zt_compressed_model_ckpt,map_location='cpu'), strict=True)
        zt_compressed_model.eval().requires_grad_(False).to(dtype=dtype)
        if args.use_xformers:
            zt_compressed_model.enable_xformers_memory_efficient_attention()
        pipe.unet = HybridModel(zt_compressed_model, pipe.unet, zt_compress_rate=args.zt_compress_rate)
        
    extra_kwargs={}
    if hasattr(args, 'control_scale'):
        extra_kwargs['controlnet_conditioning_scale']=args.control_scale

    if args.use_xformers:
        pipe.enable_xformers_memory_efficient_attention()
    pipe = pipe.to(device)
    
    if args.redit:
        redit_pipeline_cls = _PIPELINES[args.redit_pipeline]
        redit_pipeline = redit_pipeline_cls.from_pretrained(
            args.redit_pipeline_ckpt,
            torch_dtype=dtype,
            safety_checker=None,
            requires_safety_checker=False,
        )
        redit_scheduler_cls = _SCHEDULERS[args.redit_scheduler]
        redit_pipeline.scheduler = redit_scheduler_cls.from_config(redit_pipeline.scheduler.config)
        if args.use_xformers:
            redit_pipeline.enable_xformers_memory_efficient_attention()
        redit_pipeline = redit_pipeline.to(device)
        redit_pipeline.scheduler.set_timesteps(args.redit_inference_steps, device=device)

    def sample_images(prompt):
        sample_latent = pipe.inference_latent(
            prompt=prompt,
            num_inference_steps=args.infer_steps,
            guidance_scale=args.guidance_scale,
            disable_pbar=True,
            **extra_kwargs,
        )['latent']
        if args.redit:
            noisy_latent, start_step = redit_pipeline.add_noise(args.redit_inference_steps, strength=args.redit_strength, latents=sample_latent, )
            redit_latent = redit_pipeline.inference_latent(
                prompt=prompt,
                latents=noisy_latent,
                num_inference_steps=args.redit_inference_steps,
                guidance_scale=args.redit_guidance_scale,
                disable_pbar=True,
                start_step=start_step,
                **extra_kwargs,
            )['latent']
            sample_latent = redit_latent
        return pipe.latent2image(sample_latent)

    i_start, i_end = idx_se
    idx = 0
    bs=args.bs
    time_start = time.time()
    while idx < len(caps):
        if idx+bs >= len(caps):
            bs = len(caps)-idx
        prompt = caps[idx:idx+bs]
        s, e = idx+i_start, idx+bs+i_start
        
        if not hasattr(args, 'num_candidates'):
            images = sample_images(prompt).images
            for i, image in enumerate(images, start=s):
                image.save(image_dir / f'{i}.png')
        else:
            for candidate in range(args.num_candidates):
                images = sample_images(prompt).images
                for i, image in enumerate(images, start=s):
                    image.save(image_dir / f'{i}' / f'{candidate}.png')
        
        idx += bs
        logger.info(
            f'{device}: [{idx}/{len(caps)}] ; time elapased {time.time()-time_start:.3f}')

def prepare_args():
    parser = ArgumentParser()
    parser.add_argument(
        "--metric_type",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--caption_path",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--save_dir",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--devices",
        type=str,
        default=[0],
    )
    parser.add_argument(
        "--infer_steps",
        type=int,
        default=25,
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
        default=7.5,
    )
    parser.add_argument(
        "--bs",
        type=int,
        default=10,
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=58348673,
    )
    parser.add_argument(
        "--pipeline",
        type=str,
        default='alt_diffusion',
    )
    parser.add_argument(
        "--scheduler",
        type=str,
        default="dpm_solver_multi",
    )
    parser.add_argument(
        "--pipeline_ckpt",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--ckpt_path",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--num_candidates",
        type=int,
        default=24,
        help="only useful when metric type is 'CLIPSim'"
    )
    
    parser.add_argument(
        "--load_lora",
        action='store_true',
        default=False,
    )
    parser.add_argument(
        "--lora_weights_path",
        type=str,
        default=None,
    )
    
    parser.add_argument(
        "--load_controlnet",
        action='store_true',
        default=False,
    )
    parser.add_argument(
        "--controlnet_path",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--control_scale",
        type=float,
        default=1.0,
    )
    
    parser.add_argument(
        "--load_zt_compressed_model",
        action='store_true',
        default=False,
    )
    parser.add_argument(
        "--zt_compressed_model_ckpt",
        type=str,
        default=None,
    )
    parser.add_argument(
        "--zt_compress_rate",
        type=float,
        default=1.0,
    )
    
    parser.add_argument(
        "--redit",
        action='store_true',
        default=False,
    )
    parser.add_argument(
        "--redit_pipeline",
        type=str,
        default='alt_diffusion',
    )
    parser.add_argument(
        "--redit_scheduler",
        type=str,
        default='dpm_solver_multi',
    )
    parser.add_argument(
        "--redit_pipeline_ckpt",
        type=str,
    )
    parser.add_argument(
        "--redit_inference_steps",
        type=int,
        default=25,
    )
    parser.add_argument(
        "--redit_guidance_scale",
        type=int,
        default=7.5,
    )
    parser.add_argument(
        "--redit_strength",
        type=float,
        default=0.2,
    )
    
    parser.add_argument(
        "--use_xformers",
        action='store_true',
        default=False,
    )
    
    parser.add_argument(
        "--overide_timesteps",
        type=float,
        nargs='+',
        default=None, 
    )
    
    args= parser.parse_args()
    
    if isinstance(args.devices, str):
        args.devices = [int(i) for i in (args.devices).rstrip(',').split(',')]

    if args.metric_type not in ['CLIPSim']:
        del args.num_candidates
        
    if not (use_controlnet and args.load_controlnet):
        del args.control_scale
        
    return args


def save_config(args, save_path):
    config2save = vars(args)
    config2save.pop('save_dir')
    json.dump(config2save, open(save_path, 'w'))


if __name__ == "__main__":
    mp.set_start_method('spawn')
    args = prepare_args()

    # 加载 caption
    with open(args.caption_path, 'r') as f:
        captions = f.read().splitlines()
    args.num_samples = len(captions)
    set_seed(args.seed)

    # 创建需要的文件夹
    save_dir = Path(args.save_dir) / args.metric_type
    save_dir.mkdir(parents=True, exist_ok=True)
    os.system(f'cp {args.caption_path} {str(save_dir / "captions.txt")}')

    logger.add(str(save_dir / f'run.log'))
    logger.info(f'total samples {args.num_samples}')
    
    assert args.pipeline in _PIPELINES , f'available pipelines: {list(_PIPELINES.keys())}'
    assert args.scheduler in _SCHEDULERS, f'available schedulers: {list(_SCHEDULERS.keys())}'

    metric_type = args.metric_type

    logger.info(f'metric type {metric_type}')

    image_dir = save_dir / 'images'
    image_dir.mkdir(parents=True, exist_ok=True)

    main_func_args = (
        args,
        save_dir, 
        image_dir,
    )
    main_func=produce
    if metric_type in ['CLIPSim']:
        # 1. 将 caption 存储为 json形式
        logger.info(f'save captions to json')
        captions_json = {f'{i}': caption for i, caption in enumerate(captions)}
        json.dump(captions_json, open(str(save_dir / 'captions.json'), 'w'))
        # 2. 创建子文件夹
        logger.info(f'making subdirs per caption')
        for i in trange(args.num_samples):
            (image_dir / f'{i}').mkdir(parents=True, exist_ok=True)

    device_list = [f'cuda:{i}' for i in args.devices]
    num_workers = len(device_list)
    logger.info(f'use devices {args.devices} ; {num_workers} in total')

    save_config(args, str(save_dir / 'index.json'))

    groups, groups_se = devide_to_groups(captions, num_workers)
    workers = []
    for i in range(num_workers):
        p = mp.Process(target=main_func, args=(
            groups[i], groups_se[i], device_list[i],) + main_func_args)
        p.start()
        workers.append(p)
        logger.info(f'process in {device_list[i]} started')

    for p in workers:
        p.join()
        
    logger.info(f'done')
